# -*- coding: utf-8 -*-
import os
from rknn.api import RKNN
import onnx

# --- 用户配置区 ---
CONFIG = {
    # 训练好的 PyTorch 模型路径
    "PT_MODEL_PATH": "./yolov8s/best8s.pt",

    # ONNX 和 RKNN 模型的输出目录
    "OUTPUT_DIR": "yolov8s",

    # 目标平台
    "TARGET_PLATFORM": "rk3576",

    # 用于量化的数据集文件
    "QUANTIZATION_DATASET": "./dataset.txt",

    # 模型输入的尺寸
    "MODEL_INPUT_SIZE": (640, 640),
    
    # 类别数量
    "NUM_CLASSES": 3,
}
# --- 用户配置区结束 ---

def export_onnx_from_pt(pt_model_path, onnx_model_path, model_input_size):
    """将 YOLOv8 的 .pt 模型导出为 ONNX 格式"""
    print(f">>> 正在从 '{pt_model_path}' 导出 ONNX 模型...")
    if os.path.exists(onnx_model_path):
        print(f"--- ONNX 模型 '{onnx_model_path}' 已存在，跳过导出。")
        return True

    try:
        from ultralytics import YOLO
        model = YOLO(pt_model_path)
        model.export(format='onnx', imgsz=model_input_size[0], opset=12)
        print(f">>> ONNX 模型已成功导出到 '{onnx_model_path}'")
        return True
    except Exception as e:
        print(f"!!! 导出 ONNX 模型失败: {e}")
        return False

def split_yolov8_output(onnx_model_path, new_onnx_model_path, num_classes):
    """
    加载YOLOv8 ONNX模型，将其单一输出分割为box和class两个独立的输出，
    并增加Identity节点作为屏障，防止RKNN编译器进行不当优化。
    """
    print(">>> 正在分割ONNX模型的输出...")
    if os.path.exists(new_onnx_model_path):
        print(f"--- 分割后的ONNX模型 '{new_onnx_model_path}' 已存在，跳过分割。")
        return True

    try:
        model = onnx.load(onnx_model_path)
        graph = model.graph

        # 找到原始输出节点及其生产者
        original_output = graph.output[0]
        output_name = original_output.name
        
        # 从图中移除原始输出定义
        graph.output.remove(original_output)

        # 创建分割后的中间节点名称
        box_split_name = output_name + "_box_split"
        cls_split_name = output_name + "_cls_split"

        # 1. 添加一个 Split 节点来分割原始输出张量
        split_node = onnx.helper.make_node(
            'Split',
            inputs=[output_name],
            outputs=[box_split_name, cls_split_name],
            name='split_box_and_class_outputs',
            axis=1,
            split=[4, num_classes]
        )
        graph.node.append(split_node)

        # 2. 【关键修改】添加 Identity 节点作为优化屏障
        final_box_output_name = output_name + "_box"
        final_cls_output_name = output_name + "_cls"

        identity_box_node = onnx.helper.make_node(
            'Identity',
            inputs=[box_split_name],
            outputs=[final_box_output_name],
            name='identity_barrier_box'
        )
        identity_cls_node = onnx.helper.make_node(
            'Identity',
            inputs=[cls_split_name],
            outputs=[final_cls_output_name],
            name='identity_barrier_cls'
        )
        graph.node.extend([identity_box_node, identity_cls_node])

        # 3. 为最终的输出创建 ValueInfoProto 并添加到图中
        box_output_info = onnx.helper.make_tensor_value_info(
            final_box_output_name,
            onnx.TensorProto.FLOAT,
            ['batch', 4, 'predictions']
        )
        cls_output_info = onnx.helper.make_tensor_value_info(
            final_cls_output_name,
            onnx.TensorProto.FLOAT,
            ['batch', num_classes, 'predictions']
        )
        graph.output.extend([box_output_info, cls_output_info])

        # 检查并保存修改后的模型
        onnx.checker.check_model(model)
        onnx.save(model, new_onnx_model_path)
        print(f">>> 分割后的ONNX模型已成功保存到 '{new_onnx_model_path}'")
        return True

    except Exception as e:
        print(f"!!! 分割 ONNX 模型失败: {e}")
        return False

def build_rknn_from_onnx(onnx_model_path, rknn_model_path, target_platform, quant_dataset_path):
    """从 ONNX 模型构建并量化为 RKNN 模型"""
    print(f">>> 正在从 '{onnx_model_path}' 构建 RKNN 模型...")
    if os.path.exists(rknn_model_path):
        print(f"--- RKNN 模型 '{rknn_model_path}' 已存在，跳过构建。")
        return True

    rknn = RKNN(verbose=True)

    # 1. 配置 RKNN
    print("--- 1. 配置 RKNN Toolkit...")
    rknn.config(
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        target_platform=target_platform,
        quantized_dtype="w8a8",
        quantized_algorithm='normal',
        optimization_level=3
    )

    # 2. 加载 ONNX 模型
    print("--- 2. 加载 ONNX 模型...")
    ret = rknn.load_onnx(model=onnx_model_path)
    if ret != 0:
        print("!!! 加载 ONNX 模型失败!")
        rknn.release()
        return False

    # 3. 构建 RKNN 模型 (包含量化)
    print("--- 3. 构建 RKNN 模型 (此过程将进行量化，可能需要几分钟)...")
    if not os.path.exists(quant_dataset_path) or os.path.getsize(quant_dataset_path) == 0:
        print(f"!!! 错误: 量化数据集文件 '{quant_dataset_path}' 不存在或为空。")
        rknn.release()
        return False

    ret = rknn.build(do_quantization=True, dataset=quant_dataset_path)
    if ret != 0:
        print("!!! 构建 RKNN 模型失败!")
        rknn.release()
        return False

    # 4. 导出 RKNN 模型
    print("--- 4. 导出 RKNN 模型...")
    ret = rknn.export_rknn(rknn_model_path)
    if ret != 0:
        print("!!! 导出 RKNN 模型失败!")
    else:
        print(f">>> RKNN 模型已成功导出到 '{rknn_model_path}'")

    rknn.release()
    return ret == 0

def main():
    """主执行函数"""
    if not os.path.exists(CONFIG["OUTPUT_DIR"]):
        os.makedirs(CONFIG["OUTPUT_DIR"])

    base_name = os.path.splitext(os.path.basename(CONFIG["PT_MODEL_PATH"]))[0]
    onnx_path = os.path.join(CONFIG["OUTPUT_DIR"], f"{base_name}.onnx")
    split_onnx_path = os.path.join(CONFIG["OUTPUT_DIR"], f"{base_name}_split.onnx")
    rknn_path = os.path.join(CONFIG["OUTPUT_DIR"], f"{base_name}.rknn")

    # 步骤 1：从 .pt 导出原始 ONNX 模型
    if not export_onnx_from_pt(CONFIG["PT_MODEL_PATH"], onnx_path, CONFIG["MODEL_INPUT_SIZE"]):
        return

    # 步骤 2：分割 ONNX 模型的输出
    if not split_yolov8_output(onnx_path, split_onnx_path, CONFIG["NUM_CLASSES"]):
        return

    # 步骤 3：从分割后的 ONNX 模型构建 RKNN 模型
    if not build_rknn_from_onnx(split_onnx_path, rknn_path, CONFIG["TARGET_PLATFORM"], CONFIG["QUANTIZATION_DATASET"]):
        return

    print("\n>>> 模型转换流程全部完成！")

if __name__ == '__main__':
    main()


